import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np
import base64
import io
from typing import Dict, Any, Optional

class ChartGenerator:
    """图表生成器"""
    
    def __init__(self):
        # 设置 seaborn 样式
        sns.set_style("whitegrid")
        plt.rcParams['figure.figsize'] = (10, 6)
        plt.rcParams['font.size'] = 10
    
    def generate_chart(self, df: pd.DataFrame, chart_type: str, 
                      x_column: Optional[str] = None, 
                      y_column: Optional[str] = None) -> Dict[str, Any]:
        """生成图表"""
        
        try:
            plt.figure(figsize=(12, 8))
            
            if chart_type == 'histogram':
                self._generate_histogram(df, x_column)
            elif chart_type == 'boxplot':
                self._generate_boxplot(df, x_column)
            elif chart_type == 'bar':
                self._generate_bar_chart(df, x_column)
            elif chart_type == 'pie':
                self._generate_pie_chart(df, x_column)
            elif chart_type == 'scatter':
                self._generate_scatter_plot(df, x_column, y_column)
            elif chart_type == 'correlation':
                self._generate_correlation_heatmap(df)
            elif chart_type == 'line':
                self._generate_line_chart(df, x_column, y_column)
            else:
                raise ValueError(f"不支持的图表类型: {chart_type}")
            
            # 转换为 base64
            img_base64 = self._convert_to_base64()
            
            return {
                'success': True,
                'chart_data': img_base64,
                'chart_type': chart_type
            }
            
        except Exception as e:
            plt.close()
            return {
                'success': False,
                'error': str(e)
            }
    
    def _generate_histogram(self, df: pd.DataFrame, column: str):
        """生成直方图"""
        data = df[column].dropna()
        plt.hist(data, bins=30, alpha=0.7, color='skyblue', edgecolor='black')
        plt.title(f'{column} 分布直方图', fontsize=14, fontweight='bold')
        plt.xlabel(column)
        plt.ylabel('频次')
        plt.grid(True, alpha=0.3)
        
        # 添加统计信息
        mean_val = data.mean()
        plt.axvline(mean_val, color='red', linestyle='--', 
                   label=f'均值: {mean_val:.2f}')
        plt.legend()
    
    def _generate_boxplot(self, df: pd.DataFrame, column: str):
        """生成箱线图"""
        data = df[column].dropna()
        plt.boxplot(data, patch_artist=True, 
                   boxprops=dict(facecolor='lightblue', alpha=0.7))
        plt.title(f'{column} 箱线图', fontsize=14, fontweight='bold')
        plt.ylabel(column)
        plt.grid(True, alpha=0.3)
        
        # 添加统计信息文本
        q1, median, q3 = data.quantile([0.25, 0.5, 0.75])
        plt.text(1.1, median, f'中位数: {median:.2f}', 
                verticalalignment='center')
    
    def _generate_bar_chart(self, df: pd.DataFrame, column: str):
        """生成柱状图"""
        value_counts = df[column].value_counts().head(15)
        
        plt.bar(range(len(value_counts)), value_counts.values, 
               color='lightcoral', alpha=0.8)
        plt.title(f'{column} 类别分布', fontsize=14, fontweight='bold')
        plt.xlabel(column)
        plt.ylabel('计数')
        plt.xticks(range(len(value_counts)), value_counts.index, 
                  rotation=45, ha='right')
        plt.grid(True, alpha=0.3)
        
        # 添加数值标签
        for i, v in enumerate(value_counts.values):
            plt.text(i, v + max(value_counts.values) * 0.01, str(v), 
                    ha='center', va='bottom')
    
    def _generate_pie_chart(self, df: pd.DataFrame, column: str):
        """生成饼图"""
        value_counts = df[column].value_counts().head(8)
        
        # 计算其他类别
        if len(df[column].value_counts()) > 8:
            others = df[column].value_counts().iloc[8:].sum()
            value_counts['其他'] = others
        
        colors = plt.cm.Set3(np.linspace(0, 1, len(value_counts)))
        
        plt.pie(value_counts.values, labels=value_counts.index, 
               autopct='%1.1f%%', colors=colors, startangle=90)
        plt.title(f'{column} 分布饼图', fontsize=14, fontweight='bold')
        plt.axis('equal')
    
    def _generate_scatter_plot(self, df: pd.DataFrame, x_column: str, y_column: str):
        """生成散点图"""
        clean_df = df[[x_column, y_column]].dropna()
        
        plt.scatter(clean_df[x_column], clean_df[y_column], 
                   alpha=0.6, color='steelblue', s=50)
        plt.title(f'{x_column} vs {y_column}', fontsize=14, fontweight='bold')
        plt.xlabel(x_column)
        plt.ylabel(y_column)
        plt.grid(True, alpha=0.3)
        
        # 添加趋势线
        try:
            z = np.polyfit(clean_df[x_column], clean_df[y_column], 1)
            p = np.poly1d(z)
            plt.plot(clean_df[x_column], p(clean_df[x_column]), 
                    "r--", alpha=0.8, label='趋势线')
            
            # 计算相关系数
            corr = clean_df[x_column].corr(clean_df[y_column])
            plt.text(0.05, 0.95, f'相关系数: {corr:.3f}', 
                    transform=plt.gca().transAxes, 
                    bbox=dict(boxstyle="round", facecolor='wheat', alpha=0.8))
            plt.legend()
        except:
            pass
    
    def _generate_correlation_heatmap(self, df: pd.DataFrame):
        """生成相关性热力图"""
        numeric_df = df.select_dtypes(include=[np.number])
        
        if numeric_df.empty:
            raise ValueError("没有数值型数据可用于相关性分析")
        
        correlation_matrix = numeric_df.corr()
        
        mask = np.triu(np.ones_like(correlation_matrix, dtype=bool))
        
        sns.heatmap(correlation_matrix, mask=mask, annot=True, 
                   cmap='coolwarm', center=0, square=True,
                   linewidths=0.5, cbar_kws={"shrink": .8})
        plt.title('变量相关性热力图', fontsize=14, fontweight='bold')
        plt.tight_layout()
    
    def _generate_line_chart(self, df: pd.DataFrame, x_column: str, y_column: str):
        """生成折线图"""
        clean_df = df[[x_column, y_column]].dropna().sort_values(x_column)
        
        plt.plot(clean_df[x_column], clean_df[y_column], 
                marker='o', linewidth=2, markersize=4, alpha=0.8)
        plt.title(f'{y_column} 随 {x_column} 变化趋势', 
                 fontsize=14, fontweight='bold')
        plt.xlabel(x_column)
        plt.ylabel(y_column)
        plt.grid(True, alpha=0.3)
        
        # 添加数据点数量信息
        plt.text(0.02, 0.98, f'数据点: {len(clean_df)}', 
                transform=plt.gca().transAxes, verticalalignment='top',
                bbox=dict(boxstyle="round", facecolor='lightgray', alpha=0.8))
    
    def _convert_to_base64(self) -> str:
        """转换图表为 base64 字符串"""
        buffer = io.BytesIO()
        plt.tight_layout()
        plt.savefig(buffer, format='png', dpi=150, bbox_inches='tight',
                   facecolor='white', edgecolor='none')
        buffer.seek(0)
        
        image_base64 = base64.b64encode(buffer.getvalue()).decode()
        plt.close()
        
        return f"data:image/png;base64,{image_base64}"